package com.efeichong.dynamic;

import com.zaxxer.hikari.HikariDataSource;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.MutablePropertyValues;
import org.springframework.beans.factory.support.BeanDefinitionRegistry;
import org.springframework.beans.factory.support.GenericBeanDefinition;
import org.springframework.boot.context.properties.bind.Bindable;
import org.springframework.boot.context.properties.bind.Binder;
import org.springframework.boot.context.properties.source.ConfigurationPropertyName;
import org.springframework.boot.context.properties.source.ConfigurationPropertyNameAliases;
import org.springframework.boot.context.properties.source.ConfigurationPropertySource;
import org.springframework.boot.context.properties.source.MapConfigurationPropertySource;
import org.springframework.context.EnvironmentAware;
import org.springframework.context.annotation.ImportBeanDefinitionRegistrar;
import org.springframework.core.env.Environment;
import org.springframework.core.type.AnnotationMetadata;
import org.springframework.util.StringUtils;

import javax.sql.DataSource;
import java.util.HashMap;
import java.util.Map;

/**
 * @author lxk
 * @date 2022/8/28
 * @description 动态注册数据源
 */
@Slf4j
public class DynamicDataSourceRegister implements ImportBeanDefinitionRegistrar, EnvironmentAware {

    /**
     * 别名
     */
    private final static ConfigurationPropertyNameAliases aliases = new ConfigurationPropertyNameAliases();

    /**
     * 由于部分数据源配置不同，所以在此处添加别名，避免切换数据源出现某些参数无法注入的情况
     */
    static {
        aliases.addAliases("url", new String[]{"jdbc-url"});
        aliases.addAliases("username", new String[]{"user"});
    }

    //存储注册的数据源
    private Map<String, DataSource> customDataSources = new HashMap<>();
    //配置上下文（也可以理解为配置文件的获取工具）
    private Environment evn;
    //参数绑定工具 springboot2.0新推出
    private Binder binder;

    @Override
    public void registerBeanDefinitions(AnnotationMetadata annotationMetadata, BeanDefinitionRegistry beanDefinitionRegistry) {
        DataSource defaultDatasource = null;
        Map dsProperties = binder.bind("spring.datasource", Map.class).get();
        if (dsProperties.containsKey("dynamic")) {
            Map dynamicDs = (Map) dsProperties.get("dynamic");
            int index = 0;
            for (Object o : dynamicDs.entrySet()) {
                Map.Entry dsEntry = (Map.Entry) o;
                String name = (String) dsEntry.getKey();
                Map value = (Map) dsEntry.getValue();
                //获取数据源类型
                Class<? extends DataSource> clazz = getDataSourceType((String) value.get("type"));
                DataSource consumerDatasource = bind(clazz, value);
                if (index == 0) {
                    defaultDatasource = consumerDatasource;
                }
               DynamicDataSourceContextHolder.addId(name);
                customDataSources.put(name, consumerDatasource);
                log.info("注册数据源{}成功", name);
                index++;
            }
            // bean定义类
            GenericBeanDefinition define = new GenericBeanDefinition();
            // 设置bean的类型，此处DynamicRoutingDataSource是继承AbstractRoutingDataSource的实现类
            define.setBeanClass(DynamicDataSource.class);
            // 需要注入的参数
            MutablePropertyValues mpv = define.getPropertyValues();
            // 添加默认数据源，避免key不存在的情况没有数据源可用
            mpv.add("defaultTargetDataSource", defaultDatasource);
            // 添加其他数据源
            mpv.add("targetDataSources", customDataSources);
            // 将该bean注册为datasource，不使用springboot自动生成的datasource
            beanDefinitionRegistry.registerBeanDefinition("datasource", define);
            log.info("注册数据源成功，一共注册{}个数据源", customDataSources.keySet().size());
        }
    }

    /**
     * 通过字符串获取数据源class对象
     *
     * @param typeStr
     * @return
     */
    private Class<? extends DataSource> getDataSourceType(String typeStr) {
        Class<? extends DataSource> type;
        try {
            if (StringUtils.hasLength(typeStr)) {
                // 字符串不为空则通过反射获取class对象
                type = (Class<? extends DataSource>) Class.forName(typeStr);
            } else {
                // 默认为Druid数据源
                type = HikariDataSource.class;
            }
            return type;
        } catch (Exception e) {
            throw new IllegalArgumentException("can not resolve class with type: " + typeStr); //无法通过反射获取class对象的情况则抛出异常，该情况一般是写错了，所以此次抛出一个runtimeexception
        }
    }

    /**
     * 绑定参数，以下三个方法都是参考DataSourceBuilder的bind方法实现的，目的是尽量保证我们自己添加的数据源构造过程与springboot保持一致
     *
     * @param result
     * @param properties
     */
    private void bind(DataSource result, Map properties) {
        ConfigurationPropertySource source = new MapConfigurationPropertySource(properties);
        Binder binder = new Binder(new ConfigurationPropertySource[]{source.withAliases(aliases)});
        // 将参数绑定到对象
        binder.bind(ConfigurationPropertyName.EMPTY, Bindable.ofInstance(result));
    }

    private <T extends DataSource> T bind(Class<T> clazz, Map properties) {
        ConfigurationPropertySource source = new MapConfigurationPropertySource(properties);
        Binder binder = new Binder(new ConfigurationPropertySource[]{source.withAliases(aliases)});
        // 通过类型绑定参数并获得实例对象
        return binder.bind(ConfigurationPropertyName.EMPTY, Bindable.of(clazz)).get();
    }

    /**
     * @param clazz
     * @param sourcePath 参数路径，对应配置文件中的值，如: spring.datasource
     * @param <T>
     * @return
     */
    private <T extends DataSource> T bind(Class<T> clazz, String sourcePath) {
        Map properties = binder.bind(sourcePath, Map.class).get();
        return bind(clazz, properties);
    }

    /**
     * EnvironmentAware接口的实现方法，通过aware的方式注入，此处是environment对象
     *
     * @param environment
     */
    @Override
    public void setEnvironment(Environment environment) {
        log.info("开始注册数据源");
        this.evn = environment;
        // 绑定配置器
        binder = Binder.get(evn);
    }
}